"""
AI医疗问答服务
基于向量数据库回答医疗相关问题
"""
import os
import logging
import json
import requests
import time
from .document_processor import DocProcessor

# 尝试导入必要的库
try:
    import dashscope
    from dashscope.aigc.generation import Generation
    from langchain_core.messages import HumanMessage, SystemMessage
except ImportError:
    logging.warning("DashScope库未安装，请安装: pip install dashscope")

logger = logging.getLogger(__name__)

class QAService:
    """AI问答服务"""
    
    def __init__(self):
        """初始化服务"""
        self.doc_processor = DocProcessor()
        self.vectorstore = None
        self.is_initialized = False
        self.api_key = os.environ.get("DASHSCOPE_API_KEY")
    
    def initialize(self, force_rebuild=False):
        """初始化问答系统"""
        try:
            # 检查是否需要重建向量数据库
            if force_rebuild or not os.path.exists(self.doc_processor.db_path):
                logger.info("开始处理文档并创建向量数据库...")
                success = self.doc_processor.process_doc()
                if not success:
                    logger.error("处理文档失败")
                    return False
            
            # 加载向量数据库
            self.vectorstore = self.doc_processor.load_vector_db()
            if not self.vectorstore:
                logger.error("加载向量数据库失败")
                return False
            
            # 检查API密钥是否设置
            if not self.api_key:
                logger.warning("未设置DASHSCOPE_API_KEY环境变量")
            
            self.is_initialized = True
            logger.info("问答系统初始化成功")
            return True
            
        except Exception as e:
            logger.error(f"初始化问答系统出错: {str(e)}")
            return False
    
    def call_qwen_api(self, system_message, user_message):
        """调用通义千问API"""
        if not self.api_key:
            logger.error("未设置API密钥")
            return None
            
        try:
            url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            data = {
                "model": "qwen-max",
                "input": {
                    "messages": [
                        {"role": "system", "content": system_message},
                        {"role": "user", "content": user_message}
                    ]
                },
                "parameters": {
                    "temperature": 0.7,
                    "top_p": 0.8,
                    "result_format": "message"
                }
            }
            
            logger.info(f"调用通义千问API，问题：{user_message}")
            response = requests.post(url, headers=headers, json=data)
            
            if response.status_code == 200:
                result = response.json()
                logger.info(f"API响应: {result}")
                if "output" in result and "choices" in result["output"] and len(result["output"]["choices"]) > 0:
                    if "message" in result["output"]["choices"][0] and "content" in result["output"]["choices"][0]["message"]:
                        answer = result["output"]["choices"][0]["message"]["content"]
                        logger.info(f"LLM回答: {answer}")
                        return answer
                    else:
                        logger.error(f"API响应格式不正确: {result}")
                        return None
                else:
                    logger.error(f"API响应格式不正确: {result}")
                    return None
            else:
                logger.error(f"API调用失败: {response.status_code} - {response.text}")
                return None
                
        except Exception as e:
            logger.error(f"调用API出错: {str(e)}")
            return None
    
    def call_qwen_api_stream(self, system_message, user_message, callback=None):
        """流式调用通义千问API"""
        if not self.api_key:
            logger.error("未设置API密钥")
            return None
            
        try:
            url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "Accept": "text/event-stream"  # 使用SSE流式传输
            }
            data = {
                "model": "qwen-max",
                "input": {
                    "messages": [
                        {"role": "system", "content": system_message},
                        {"role": "user", "content": user_message}
                    ]
                },
                "parameters": {
                    "temperature": 0.7,
                    "top_p": 0.8,
                    "result_format": "message",
                    "incremental_output": True  # 启用增量输出
                }
            }
            
            logger.info(f"流式调用通义千问API，问题：{user_message}")
            
            # 使用stream=True参数进行流式请求
            response = requests.post(url, headers=headers, json=data, stream=True)
            
            if response.status_code == 200:
                logger.info("API响应状态码正常，开始处理流式响应")
                full_answer = ""
                chunk_count = 0
                
                for line in response.iter_lines():
                    if line:
                        # 解析SSE格式的数据
                        line = line.decode('utf-8')
                        if line.startswith('data:'):
                            data_str = line[5:].strip()
                            if data_str:
                                try:
                                    data_json = json.loads(data_str)
                                    if "output" in data_json and "choices" in data_json["output"] and len(data_json["output"]["choices"]) > 0:
                                        choice = data_json["output"]["choices"][0]
                                        if "message" in choice and "content" in choice["message"]:
                                            chunk = choice["message"]["content"]
                                            if chunk:  # 确保chunk不为空
                                                chunk_count += 1
                                                full_answer += chunk
                                                
                                                logger.debug(f"收到第{chunk_count}个流式片段: {chunk[:20]}..." if len(chunk) > 20 else chunk)
                                                
                                                if callback:
                                                    logger.debug(f"调用回调函数处理第{chunk_count}个流式片段")
                                                    callback(chunk)
                                        else:
                                            logger.warning(f"响应中没有message.content字段: {choice}")
                                    else:
                                        logger.warning(f"响应格式不符合预期: {data_json}")
                                except json.JSONDecodeError:
                                    logger.error(f"解析JSON出错: {data_str}")
                
                logger.info(f"流式回答完成，总长度: {len(full_answer)}，片段数: {chunk_count}")
                
                # 如果没有收到任何内容，返回默认回答
                if not full_answer:
                    default_answer = "抱歉，我无法回答您的问题。请尝试重新提问或联系客服。"
                    logger.warning("流式回答未返回任何内容，使用默认回答")
                    if callback:
                        callback(default_answer)
                    return default_answer
                    
                return full_answer
            else:
                logger.error(f"流式API调用失败: {response.status_code} - {response.text}")
                error_message = "抱歉，AI服务暂时不可用，请稍后再试。"
                if callback:
                    callback(error_message)
                return error_message
                
        except Exception as e:
            logger.error(f"流式调用API出错: {str(e)}")
            error_message = "抱歉，发生了错误，请稍后再试。"
            if callback:
                callback(error_message)
            return error_message
            
    def answer(self, question, stream_callback=None):
        """回答问题"""
        if not question:
            return "请输入您的问题"
        
        # 如果尚未初始化，先初始化并强制重建向量数据库
        if not self.is_initialized:
            success = self.initialize(force_rebuild=True)
            if not success:
                return "我暂时不知道快让朱林哥哥来教教我吧~~~"
        
        try:
            # 手动实现问答过程
            if self.vectorstore:
                # 从向量数据库中检索相关文档
                docs = self.vectorstore.similarity_search(question, k=3)
                logger.info(f"检索到{len(docs)}个相关文档")
                
                # 提取文档内容
                context = "\n\n".join([doc.page_content for doc in docs])
                
                # 构建提示
                system_message = f"你是一个医疗AI助手，基于以下医疗信息回答用户问题。如果无法从提供的信息中找到答案，请诚实地说你不知道。\n\n医疗信息：\n{context}"
                
                # 调用API
                if stream_callback:
                    # 使用流式API
                    answer = self.call_qwen_api_stream(system_message, question, stream_callback)
                else:
                    # 使用普通API
                    answer = self.call_qwen_api(system_message, question)
                    
                if not answer:
                    answer = "我暂时不知道快让朱林哥哥来教教我吧~~~"
            else:
                # 如果向量数据库创建失败，使用简单回复
                logger.warning("向量数据库未创建，使用简单回复")
                answer = "我是医疗AI助手，可以回答您的医疗问题。但目前我的问答系统还在完善中，请稍后再试。"
            
            # 如果没有找到回答
            if not answer or not answer.strip():
                return "我暂时不知道快让朱林哥哥来教教我吧~~~"
            
            return answer
            
        except Exception as e:
            logger.error(f"回答问题出错: {str(e)}")
            return "我暂时不知道快让朱林哥哥来教教我吧~~~"

# 创建全局问答服务实例
qa_service = QAService() 